/* process_info.c
 * Process information (pid, process name)
 *
 * $Id$
 *
 * Wireshark - Network traffic analyzer
 * By Bogdan Harjoc <harjoc@gmail.com>
 * Copyright 1998 Gerald Combs
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#ifdef _WIN32

#include "config.h"

#include <glib.h>
#include <stdio.h>
#include <stdlib.h>

#include <windows.h>
#include <winsock.h>
#include <Iphlpapi.h>
#include <Psapi.h>

#include "process_info.h"

#include <epan/ipproto.h>

typedef struct netstat {
    struct netstat* next;
    guint loc_ip;
    guint rem_ip;
    guint16 loc_port;
    guint16 rem_port;
    guint16 proto;
    guint pid;
} netstat_t;

typedef struct procinfo {
    struct procinfo* next;
    guint pid;
    char name[1];
} procinfo_t;

#define HASH_SIZE 971

static procinfo_t* process_hash[HASH_SIZE] = { 0 };
static netstat_t* addr_hash[HASH_SIZE] = { 0 };
static netstat_t* pair_hash[HASH_SIZE] = { 0 };

/**
 * 先查询PID是否已经在哈希表process_hash中，如果已存在则直接返回。
 * 否则通过ID打开进程，查询它的进程名，并存放到哈希表内。
 */
static procinfo_t* lookup_pid(guint pid)
{
    procinfo_t* e;
    HANDLE hproc;
    static wchar_t process_name[MAX_PATH] = { 0 };
    guint len;

    // 在哈希表内查询，注意循环的作用是查询具有相同哈希值的所有procinfo_t结构
    for (e = process_hash[pid % HASH_SIZE]; e; e = e->next)
        if (e->pid == pid)
            return e;

    // 可能这里需要提权？
    hproc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
    if (!hproc)
        return NULL;

    len = GetProcessImageFileName(hproc, process_name, MAX_PATH);
    CloseHandle(hproc);
    if (len == 0)
        return NULL;

    e = g_malloc(offsetof(procinfo_t, name) + (len + 1));
    if (!e)
        return NULL;

    /* must convert properly using WideCharToMultiByte here */
    _snprintf(e->name, len + 1, "%.*S", len, process_name);
    e->name[len] = 0;

    e->pid = pid;
    e->next = process_hash[pid % HASH_SIZE];
    process_hash[pid % HASH_SIZE] = e;

    return e;
}

static guint lookup_addrs(guint16 proto, guint src_ip, guint16 src_port, guint dst_ip, guint16 dst_port)
{
    guint h_local = src_ip + src_port;
    guint h_remote = dst_ip + dst_port;
    guint h_pair = h_local + h_remote;
    netstat_t* e;

    /* src:dst matches src -> dst  or  dst -> src (established/connecting/closed/...) */
    for (e = pair_hash[h_pair % HASH_SIZE]; e; e = e->next)
        if ((e->loc_ip == src_ip && e->rem_ip == dst_ip && e->loc_port == src_port && e->rem_port == dst_port) ||
            (e->loc_ip == dst_ip && e->rem_ip == src_ip && e->loc_port == dst_port && e->rem_port == src_port))
            if (e->proto == proto)
                return e->pid;

    /* src matches src -> 0.0.0.0:0 (listening) */
    for (e = addr_hash[h_local % HASH_SIZE]; e; e = e->next)
        if (e->loc_ip == src_ip && e->loc_port == src_port && e->rem_ip == 0 && e->rem_port == 0)
            if (e->proto == proto)
                return e->pid;

    /* dst matches dst -> 0.0.0.0:0 (listening) */
    for (e = addr_hash[h_remote % HASH_SIZE]; e; e = e->next)
        if (e->loc_ip == dst_ip && e->loc_port == dst_port && e->rem_ip == 0 && e->rem_port == 0)
            if (e->proto == proto)
                return e->pid;

    /* src matches src -> something  or  something -> src (netstat and winpcap see different IP addrs) */
    for (e = addr_hash[h_local % HASH_SIZE]; e; e = e->next)
        if ((e->loc_ip == src_ip && e->loc_port == src_port && e->rem_port == dst_port) ||
            (e->rem_ip == src_ip && e->rem_port == src_port && e->loc_port == dst_port))
            if (e->proto == proto)
                return e->pid;

    /* dst matches dst -> something  or  something -> dst (netstat and winpcap see different IP addrs) */
    for (e = addr_hash[h_remote % HASH_SIZE]; e; e = e->next)
        if ((e->loc_ip == dst_ip && e->loc_port == dst_port && e->rem_port == src_port) ||
            (e->rem_ip == dst_ip && e->rem_port == dst_port && e->loc_port == src_port))
            if (e->proto == proto)
                return e->pid;

    return 0;
}

static void netstat_update()
{
    guint buf_size = 0;
    guint h;
    char* buf = NULL;
    //netstat_t* e = NULL;
    int ret;
    unsigned i;
    MIB_TCPTABLE_OWNER_PID* mibTable;

    /**
     * 清除addr_hash和pair_hash两个哈希表的值。
     */
    for (h = 0; h < HASH_SIZE; h++) {
        while (addr_hash[h]) {
            netstat_t* head = addr_hash[h];
            addr_hash[h] = addr_hash[h]->next;
            g_free(head);
        }
        while (pair_hash[h]) {
            netstat_t* head = pair_hash[h];
            pair_hash[h] = pair_hash[h]->next;
            g_free(head);
        }
    }

    /* 执行netstat */
    ret = GetExtendedTcpTable(NULL, &buf_size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
    if (ret == ERROR_INSUFFICIENT_BUFFER) {
        buf = g_malloc(buf_size);
        if (!buf) return;
        ret = GetExtendedTcpTable(buf, &buf_size, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
    }
    if (ret != NO_ERROR) {
        g_free(buf);
        return;
    }

    mibTable = (MIB_TCPTABLE_OWNER_PID*)buf;
    for (i = 0; i < mibTable->dwNumEntries; i++) {
        guint src_ip = _byteswap_ulong(mibTable->table[i].dwLocalAddr);
        guint dst_ip = _byteswap_ulong(mibTable->table[i].dwRemoteAddr);
        guint src_port = _byteswap_ushort((unsigned short)mibTable->table[i].dwLocalPort);
        guint dst_port = _byteswap_ushort((unsigned short)mibTable->table[i].dwRemotePort);
        guint pid = mibTable->table[i].dwOwningPid;

        guint h_local = src_ip + src_port;
        guint h_remote = dst_ip + dst_port;
        guint h_pair = h_local + h_remote;

        netstat_t* e;

        /* 添加以下项目到哈希表中：
           - src->dst and dst->src to pair_hash[src:dst]  -- note that hash(src:dst) == hash(dst:src)
           - src->dst and dst->src to addr_hash[src]
           - src->dst and dst->src to addr_hash[dst]

           dst:src is also added in order to find packets coming in both directions
           netstat_t带有next指针，用它可以将相同哈希值的元素串在一起。
        */

        //
        // 要假设数据包为出去的包
        // netstat_t.local ip = IP包的source ip
        // netstat_t.remote ip = IP包的destination ip
        // 同时，让addr_hash[src_ip + src_port] = addr_hash[dst_ip + dst_port]
        //
        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = src_ip;
        e->rem_ip = dst_ip;
        e->loc_port = src_port;
        e->rem_port = dst_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = addr_hash[h_local % HASH_SIZE];
        addr_hash[h_local % HASH_SIZE] = e;

        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = src_ip;
        e->rem_ip = dst_ip;
        e->loc_port = src_port;
        e->rem_port = dst_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = addr_hash[h_remote % HASH_SIZE];
        addr_hash[h_remote % HASH_SIZE] = e;

        //
        // 要假设数据包是回来的包
        // netstat_t.local ip = IP数据包.destination IP.
        // netstat_t.remote ip = IP数据包.source IP
        // 同时，让addr_hash[src_ip + src_port] = addr_hash[dst_ip + dst_port]
        //
        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = dst_ip;
        e->rem_ip = src_ip;
        e->loc_port = dst_port;
        e->rem_port = src_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = addr_hash[h_local % HASH_SIZE];
        addr_hash[h_local % HASH_SIZE] = e;

        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = dst_ip;
        e->rem_ip = src_ip;
        e->loc_port = dst_port;
        e->rem_port = src_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = addr_hash[h_remote % HASH_SIZE];
        addr_hash[h_remote % HASH_SIZE] = e;

        //
        // pair_hash[src_ip + src_port + dst_ip + dst_port] = e;
        //
        // 假设是出去方向，netstat_t.local ip = IP包的source ip。
        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = src_ip; 
        e->rem_ip = dst_ip;
        e->loc_port = src_port;
        e->rem_port = dst_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = pair_hash[h_pair % HASH_SIZE];
        pair_hash[h_pair % HASH_SIZE] = e;

        // 假设是回来方向，netstat_t.local ip = IP数据包.destination IP.。
        e = g_malloc(sizeof(netstat_t));
        if (!e)
            break;
        e->loc_ip = dst_ip;
        e->rem_ip = src_ip;
        e->loc_port = dst_port;
        e->rem_port = src_port;
        e->proto = IP_PROTO_TCP;
        e->pid = pid;
        e->next = pair_hash[h_pair % HASH_SIZE];
        pair_hash[h_pair % HASH_SIZE] = e;
    }

    g_free(buf);
}

const char* process_info_lookup(tvbuff_t* tvb)
{
    static char buf[64];
    guint ip_vhl = tvb_get_guint8(tvb, 0);
    //guint ip_len = tvb_get_ntohs(tvb, 2);
    guint8 ip_proto = tvb_get_guint8(tvb, 9);
    guint proto_ofs;
    guint src_ip = tvb_get_ntohl(tvb, 12);
    guint dst_ip = tvb_get_ntohl(tvb, 16);
    guint16 src_port;
    guint16 dst_port;
    guint pid;
    procinfo_t* info;

    if (ip_vhl >> 4 != 4)
        return NULL;
    if (ip_proto != IP_PROTO_TCP /* && ip_proto != IP_PROTO_UDP */)
        return NULL;

    proto_ofs = (ip_vhl & 0xf) * 4;

    src_port = tvb_get_ntohs(tvb, proto_ofs + 0);
    dst_port = tvb_get_ntohs(tvb, proto_ofs + 2);

    pid = lookup_addrs(ip_proto, src_ip, src_port, dst_ip, dst_port);
    if (!pid) {
        netstat_update();
        pid = lookup_addrs(ip_proto, src_ip, src_port, dst_ip, dst_port);
        if (!pid)
            return NULL;
    }

    info = lookup_pid(pid);
    if (!info)
        return NULL;
    return info->name;
}

#else /* _WIN32 */

const char* process_info_lookup(tvbuff_t* tvb) { return NULL; }

#endif /* _WIN32 */
